from configparser import Interpolation
import numpy as np
import sophon.sail as sail
import cv2
import logging
import os

from utils import nms, xywh2xyxy, xyxy2xywh
import time

class Model:
    def __init__(self, bmodel, conf_thres, iou_thres, tpu_id=0):
        self.tpu_id = tpu_id
        self.bmodel = bmodel
        self.inp_size = None # set value after parsing model
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        self.set_engine()

    def set_engine(self):
        logging.info('parsing %s' % self.bmodel)
        engine = sail.Engine(self.tpu_id)
        engine.load(self.bmodel)
        
        graph_name = engine.get_graph_names()[0]
        inp_names = engine.get_input_names(graph_name)
        out_names = engine.get_output_names(graph_name)
        if len(inp_names) != 1:
            logging.error('len(inp_names) != 1')
            os._exit(-1)
        logging.info('found %d outputs: %s' % (len(out_names), out_names))

        inp_shape = engine.get_input_shape(graph_name, inp_names[0])
        if len(inp_shape) != 4:
            logging.error('len(inp_shape) != 4')
            os._exit(-1)
        self.inp_size = inp_shape[2:4]
        
        handle = engine.get_handle()
        engine.set_io_mode(graph_name, sail.IOMode.SYSO)
        self.graph_name = graph_name
        self.engine = engine
        self.engine_handle = handle
        self.inp_name = inp_names[0]
        self.out_name = out_names[-1]
        self.inp_dict = {inp_names[0]: None}
        self.out_dict = {}
        for name in out_names:
            shape = engine.get_output_shape(graph_name, name)
            dtype = engine.get_output_dtype(graph_name, name);
            self.out_dict[name] = sail.Tensor(handle, shape, dtype, True, True)
    
    def preprocess(self, img, s=None):
        in_h, in_w = self.inp_size

        h, w = img.shape[:2]

        if s is None:
            s = min(in_h / h, in_w / w)
        img = cv2.resize(img, (int(s*w), int(s*h)))
        img = img.astype(np.float32)
        
        pad_img = np.zeros([int(in_h), int(in_w), 3], dtype=np.float32)
        pad_img[:img.shape[0], :img.shape[1], :] = img
        
        pad_img = np.transpose(pad_img, [2, 0, 1])[None]
        pad_img = pad_img.astype(np.float32) / 255.0
        pad_img = np.ascontiguousarray(pad_img)
        return pad_img, s
    
    def postprocess(self, preds):
        if len(preds.shape)!=3 or preds.shape[0] != 1:
            logging.error("len(preds.shape)!=3 or preds.shape[0] != 1")
            os._exit(-1)
        preds = preds[0]
        confs = preds[:, 4]
        boxes = preds[:, :4]
        
        conf_mask = confs > self.conf_thres
        boxes = boxes[conf_mask]
        confs = confs[conf_mask]
        boxes = xywh2xyxy(boxes)
        keep = nms(boxes, confs, self.iou_thres)
        boxes = boxes[keep]
        confs = confs[keep]
        return boxes, confs

    def inference(self, img):
        img_t = sail.Tensor(self.engine_handle, img)
        self.inp_dict[self.inp_name] = img_t
        
        self.engine.process(self.graph_name, self.inp_dict, self.out_dict)
        
        real_output_shape1 = self.engine.get_output_shape(self.graph_name, self.out_name)
        out = self.out_dict[self.out_name].asnumpy(real_output_shape1)
        return out
    
    def __call__(self, img, scale=None):
        img_pre, s = self.preprocess(img, scale)
        out = self.inference(img_pre)
        boxes, confs = self.postprocess(out)
        boxes = boxes / s
        return boxes, confs

class Detector:
    def __init__(self, bmodel, conf_thres, iou_thres, min_box=(0, 0), extra_scales=(0.5, 0.25, 0.0625)):
        self.model = Model(bmodel, conf_thres, iou_thres)
        self.iou_thres = iou_thres
        self.min_box = min_box
        self.extra_scales = extra_scales
        #warmup
        image = np.random.randn(2160, 3840, 3) * 255
        image = image.astype(np.uint8)
        self(image)
    
    def set_conf_thres(self, thres):
        self.model.conf_thres = thres 
    
    def __call__(self, inp_image, margins=None):
        # margins [left, up, right, down]
        inp_h, inp_w, inp_c = inp_image.shape
        if (inp_h, inp_w, inp_c) != (2160, 3840, 3):
            logging.warn('image.shape is (%d, %d, %d) != (2160, 3840, 3)' % (inp_h, inp_w, inp_c))
        offsets = np.array([
            [0, 0],
            [0, inp_w//2],
            [inp_h//2, 0],
            [inp_h//2, inp_w//2]
        ])
        images = [inp_image[dy:dy+inp_h//2, dx:dx+inp_w//2] for dy,dx in offsets]
        boxes, confs = [], []
        for i, img in enumerate(images):
            _b, _c = self.model(img)
            dxdy = np.array([offsets[i][1], offsets[i][0]])
            dxdy = np.concatenate([dxdy, dxdy])
            boxes.append(_b + dxdy)
            confs.append(_c)
        
        for scale in self.extra_scales:
            img = cv2.resize(inp_image, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
            _b, _c = self.model(img, scale)
            _b = _b / scale
            boxes.append(_b)
            confs.append(_c)
        boxes = np.concatenate(boxes)
        confs = np.concatenate(confs)
        keep = nms(boxes, confs, self.iou_thres)
        boxes = boxes[keep]
        confs = confs[keep]

        if margins is not None:
            m = boxes[:, 0] > margins[0]
            m &= boxes[:, 1] > margins[1]
            m &= boxes[:, 2] < inp_w - margins[2]
            m &= boxes[:, 3] < inp_h - margins[3]
            boxes = boxes[m]
            confs = confs[m]
        
        # filter small boxes
        xywh = xyxy2xywh(boxes)
        m = (xywh[:, 2] > self.min_box[0]) & (xywh[:, 3] > self.min_box[0])
        boxes = boxes[m]
        confs = confs[m]
        return boxes, confs
